# -*- coding:utf-8 -*-
# @FileName :carla_env.py
# @Time :2024/6/15 下午4:37
# @Author :ShengYe
# @Des :
import math

import carla
import numpy as np

from my_controller.controller import Longitudinal_PID_controller
from my_code.my_utils import generate_dlc_path
from planner import planning_utils


class CarlaEnv:
    """
    模拟Carla自动驾驶环境。

    Attributes:
        client (carla.Client): Carla客户端。
        world (carla.World): Carla仿真世界。
        amap (carla.Map): Carla地图。
        ego_spawn_point (carla.Transform): 自动驾驶车辆的生成点。
        ego_vehicle (carla.Vehicle): 自动驾驶车辆。
        Lon_control (Longitudinal_PID_controller): 纵向控制器。
        _target_path (list): 自动驾驶路径。
        _max_throttle, _max_brake, _max_steer (float): 车辆控制参数。
        min_steer (float): 最小转向角。
        x_pre, y_pre (float): 上一个位置坐标。
        _waypoints (list): 车道中心线上的路径点。
        min_index (int): 最小距离路径点的索引。

    Methods:
        step(action, target_speed): 执行一步仿真。
        reset(): 重置仿真环境。
        clean(): 清理环境资源。
        _get_reward(): 计算当前状态的奖励。
        _set_spectator(): 设置观察者视角。
        _get_vehicle_control(action, target_speed): 获取车辆控制命令。
        _spawn_ego_vehicle(): 生成自动驾驶车辆。
        _get_pathway(): 获取车辆行驶路径。
        _get_state(): 获取当前状态。
        _get_min_index(x,        y): 计算距离最近路径点的索引。
        _terminal(): 判断是否达到终止条件。
        _get_distance(): 获取当前车辆到路径终点的距离。
    """

    def __init__(self):
        self.client = carla.Client('localhost', 2000)
        self.client.set_timeout(10.0)
        self.world = self.client.load_world('town05')
        self.amap = self.world.get_map()

        self.ego_spawn_point = self.amap.get_spawn_points()[7]

        self._target_path = self._get_pathway()
        self.ego_vehicle = self._spawn_ego_vehicle()

        self._max_throttle = 0.75
        self._max_brake = 0.3
        self._max_steer = 1
        self.min_steer = -1

        self.Lon_control = Longitudinal_PID_controller(self.ego_vehicle)

        # 定义状态空间维度和动作空间维度
        self.state_dim = 6
        self.action_dim = 1  # 例如：用于转向的横向控制角度
        self.action_range = 0.3  # 例如：横向控制角度的范围为 [-1, 1]

    def step(self, action, target_speed):
        """
        执行一步仿真。

        Args:
            action (numpy.ndarray): 当前动作。
            target_speed (float): 目标速度。

        Returns:
            tuple: (state, reward, done) 当前状态、奖励和是否终止的标志。
        """
        control = self._get_vehicle_control(action[0].item(), target_speed)
        self.ego_vehicle.apply_control(control)
        self.world.tick()
        self._set_spectator()
        return self._get_state(), self._get_reward(), self._terminal()

    def reset(self):
        """
        重置仿真环境。

        Returns:
            numpy.ndarray: 初始状态。
        """
        self.ego_vehicle.destroy()  # 销毁当前车辆
        self.ego_vehicle = self._spawn_ego_vehicle()  # 重新生成车辆
        self.Lon_control = Longitudinal_PID_controller(self.ego_vehicle)  # 重新初始化纵向控制器
        self.min_index = 0  # 重置路径索引
        self._set_spectator()  # 重置观察者视角
        return self._get_state()  # 返回初始状态

    def clean(self):
        """
        清理环境资源。
        """
        if self.ego_vehicle is not None:
            self.ego_vehicle.destroy()
        settings = self.world.get_settings()
        settings.synchronous_mode = False  # 禁用同步模式
        settings.fixed_delta_seconds = None

    def _get_reward(self):
        """
        计算当前状态的奖励。

        Returns:
            float: 当前状态的奖励。
        """
        ts = 0.01
        x, y, fi, V_y, fi_dao, V_x = self._get_state()
        x += V_x * ts * math.cos(fi) - V_y * ts * math.sin(fi)
        y += V_y * ts * math.cos(fi) + V_x * ts * math.sin(fi)
        fi += fi_dao * ts
        self.x_pre, self.y_pre = x, y

        min_index = self._get_min_index(x, y)
        self.min_index = min_index

        n_v = np.array([-math.sin(self._target_path[min_index][2]), math.cos(self._target_path[min_index][2])])
        d_v = np.array([x - self._target_path[min_index][0], y - self._target_path[min_index][1]])
        e_d = np.dot(n_v, d_v)
        return -e_d

    def _set_spectator(self):
        """
        设置观察者视角。
        """
        self.world.get_spectator().set_transform(
            carla.Transform(self.ego_vehicle.get_transform().location +
                            carla.Location(z=50),
                            carla.Rotation(pitch=-90))
        )

    def _get_vehicle_control(self, action, target_speed):
        """
        获取车辆控制命令。

        Args:
            action (float): 转向控制角度。
            target_speed (float): 目标速度。

        Returns:
            carla.VehicleControl: 车辆控制命令。
        """
        control = carla.VehicleControl()

        control.steer = max(self.min_steer, min(self._max_steer, action))

        control.throttle = min(self._max_throttle, max(0, self.Lon_control.PID_control(target_speed)))
        control.brake = max(0, -self.Lon_control.PID_control(target_speed))
        control.hand_brake = False
        control.manual_gear_shift = False
        control.gear = 1
        return control

    def _spawn_ego_vehicle(self):
        """
        生成自动驾驶车辆。

        Returns:
            carla.Vehicle: 生成的自动驾驶车辆。
        """
        model3_bp = self.world.get_blueprint_library().find('vehicle.tesla.model3')
        model3_bp.set_attribute('color', '255,88,0')
        model3_actor = self.world.spawn_actor(model3_bp, self.ego_spawn_point)

        #print("【车辆已经生成】transform为", self.ego_spawn_point)
        physics_control = carla.VehiclePhysicsControl()
        physics_control.mass = 1412
        model3_actor.apply_physics_control(physics_control)
        return model3_actor

    def _get_pathway(self):
        """
        获取车辆行驶路径。

        Returns:
            list: 车辆行驶路径。
        """
        self._waypoints = generate_dlc_path.get_waypoints_2(self.ego_spawn_point)
        return planning_utils.waypoint_list_2_target_path(self._waypoints)

    def _get_state(self):
        """
        获取当前状态。

        Returns:
            tuple: 当前状态 (x, y, fi, Vy, fi_dao, Vx)。
        """
        vehicle_loc = self.ego_vehicle.get_location()
        x, y = vehicle_loc.x, vehicle_loc.y
        fi = self.ego_vehicle.get_transform().rotation.yaw * (math.pi / 180)
        V = self.ego_vehicle.get_velocity()
        v_vec = [V.x, V.y, V.z]
        forward_vec = self.ego_vehicle.get_transform().get_forward_vector()
        Vx = max(np.dot(v_vec, [forward_vec.x, forward_vec.y, forward_vec.z]), 0.005)
        right_vec = self.ego_vehicle.get_transform().get_right_vector()
        Vy = np.dot(v_vec, [right_vec.x, right_vec.y, right_vec.z])
        fi_dao = self.ego_vehicle.get_angular_velocity().z * (math.pi / 180)
        return x, y, fi, Vy, fi_dao, Vx

    def _get_min_index(self, x, y):
        """
        计算距离最近路径点的索引。

        Args:
            x (float): 当前位置x坐标。
            y (float): 当前位置y坐标。

        Returns:
            int: 最近路径点的索引。
        """
        min_d = float('inf')
        min_index = self.min_index
        for i in range(self.min_index, min(self.min_index + 50, len(self._target_path))):
            d = (self._target_path[i][0] - x) ** 2 + (self._target_path[i][1] - y) ** 2
            if d < min_d:
                min_d = d
                min_index = i
        return min_index

    def _terminal(self):
        """
        判断是否达到终止条件。

        Returns:
            bool: 是否达到终止条件。
        """
        return self._get_distance() < 2

    def _get_distance(self):
        """
        获取当前车辆到路径终点的距离。

        Returns:
            float: 当前车辆到路径终点的距离。
        """
        location = self.ego_vehicle.get_transform().location
        return location.distance(self._waypoints[-1][0].transform.location)

